package com.hatricks.is.base.cookie;

import cn.hutool.core.util.ReflectUtil;
import cn.hutool.http.Header;
import lombok.extern.slf4j.Slf4j;
import org.apache.tomcat.util.buf.MessageBytes;
import org.apache.tomcat.util.http.MimeHeaders;

import javax.servlet.ServletResponse;
import javax.servlet.http.HttpServletResponse;
import java.util.HashSet;
import java.util.Set;


/**
 * @packageName: com.chinasofti.broker.bpm.common.filter
 * @className: CookieSameSiteUtils
 * @author: yanyuhui yanyuhui@chinasofti.com
 * @date: 2022/3/31 17:46
 * @description:
 */
@Slf4j
public class CookieSameSiteUtils {

    public static void addSameSite(HttpServletResponse response) {
        MimeHeaders mimeHeaders = getMimeHeaders(response);
        if (mimeHeaders == null) {
            return;
        }
        Set<String> cookieValues = new HashSet<>();
        for (int i = 0, size = mimeHeaders.size(); i < size; i++) {
            if (mimeHeaders.getName(i).equalsIgnoreCase(Header.SET_COOKIE.getValue())) {
                cookieValues.add(mimeHeaders.getValue(i).getString().concat(";SameSite=").concat("None"));
            }
        }
        mimeHeaders.removeHeader(Header.SET_COOKIE.getValue());
        for (String messageByte : cookieValues) {
            MessageBytes messageBytes = mimeHeaders.addValue(Header.SET_COOKIE.getValue());
            messageBytes.setString(messageByte);
        }
    }

    private static MimeHeaders getMimeHeaders(ServletResponse response) {
        Object res = getResponse(response);
        if (res != null) {
            return (MimeHeaders) ReflectUtil.getFieldValue(res, "headers");
        }
        return null;
    }

    private static Object getResponse(Object obj) {
        Object res = ReflectUtil.getFieldValue(obj, "response");
        if (res != null) {
            return getResponse(res);
        } else {
            return ReflectUtil.getFieldValue(obj, "coyoteResponse");
        }
    }


}
